import torch
from vector_quantize_pytorch import FSQ


class FSQWrapper(torch.nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.fsq = FSQ(
            levels=[3, 3, 3, 3, 3, 3, 3, 3],
            dim=dim,
            preserve_symmetry=True
        )

    def forward(self, x):
        # 明确返回 FSQ 的输出（部分版本的 FSQ 可能返回元组，需根据实际输出调整）
        out = self.fsq(x)
        # 若 FSQ 返回 (quantized, indices)，可按需选择输出（例如只导出量化特征）
        # 示例：return out[0]  # 仅返回量化后的特征
        return out


# 实例化模型（dim=1280 需与实际输入的最后一维匹配）
model = FSQWrapper(dim=1280)
model.eval()

# 测试输入：形状为 (batch_size=1, seq_len=100, dim=1280)
# 注意：最后一维必须固定为 1280（与 FSQ 的 dim 参数一致）
dummy_input = torch.randn(1, 100, 1280)

# 导出 ONNX 模型
torch.onnx.export(
    model,
    args=(dummy_input,),  # 输入参数（元组形式）
    f="fsq.onnx",
    opset_version=14,  # 提升 opset 版本以更好支持动态形状和新算子
    dynamic_axes={
        # 关键：输入名称需与模型实际输入名称一致（默认是 "x"，可通过 verbose=True 确认）
        "x": {0: "batch_size", 1: "seq_len"},  # 第0维（batch）和第1维（seq）动态
        # 输出名称需与模型实际输出名称一致（可通过 verbose=True 查看）
        "output": {0: "batch_size", 1: "seq_len"}  # 输出轴与输入轴对应
    },
    input_names=["x"],  # 显式指定输入名称（与 dynamic_axes 对应）
    output_names=["output"],  # 显式指定输出名称（与 dynamic_axes 对应）
    verbose=False  # 调试时设为 True，可查看算子细节和输入输出名称
)